import cv2
from ultralytics import YOLO


class Helmet_Detector:
    def __init__(self, model_path='best.pt'):
        # 加载自定义训练的YOLO模型
        self.model = YOLO(model_path)
        # self.classes = {0: 'helmet', 1: 'people',2:'carry objects',3:'car2',4:'hat',5:'no-helmet',6:'person'}
        self.classes = {0:'helmet', 1: 'person'}
    def detect(self, image_path):
        # 使用模型进行预测
        results = self.model.predict(source=image_path)

        # 获取检测结果
        result = results[0]
        # 绘制检测结果
        annotated_frame = result.plot()

        # 保存检测后的图像
        output_path = image_path.replace('.jpg', '_detected.jpg').replace('.png', '_detected.png')
        cv2.imwrite(output_path, annotated_frame)

        # 提取检测结果信息
        detections = []
        for box in result.boxes:
            class_id = box.cls.item()
            confidence = box.conf.item()
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            print(class_id, confidence, x1, y1, x2, y2)
            # 只保留安全帽和人员的检测结果
            # if class_id in [0, 1] and confidence > 0.5:
            detections.append({
                'class': self.classes[class_id],
                'confidence': confidence,
                'box': [x1, y1, x2, y2]
            })

        return {
            'detected_image': output_path,
            'objects': detections
        }